package com.tencent.demo.mp4writer;


import android.annotation.TargetApi;
import android.media.MediaCodec.BufferInfo;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.os.Build;
import android.util.Log;

import java.io.IOException;
import java.nio.ByteBuffer;

@TargetApi(Build.VERSION_CODES.JELLY_BEAN_MR2) 
public class MediaMP4Writer{

	public static final int TRACK_VIDEO = 0;
	public static final int TRACK_AUDIO = 1;
	private final String LOG_TAG = "MediaMP4Writer";
	
	private String  mPath = null;
	
	public boolean mStarted 		= false;
	private int 	mAudioTrackIndex = -1;
	private int 	mVideoTrackIndex = -1;
	
	MediaMuxer 	mMuxer 		 = null;
	
	public MediaMP4Writer(String path) {
	    
	    Log.i(LOG_TAG, "[MediaMP4Writer] Enter MediaMP4Writer");
	    
	    if (path == null) {
	        throw new IllegalArgumentException("path must not be null");
	    }
	    
	    mPath = path;
	    Log.i(LOG_TAG, "[MediaMP4Writer] Leave MediaMP4Writer mPath = " + mPath);
	}

	public void init() throws IOException {
		
		Log.d(LOG_TAG, "[init] Enter init: path = " + mPath);
		
        // Create a MediaMuxer.  We can't add the video track and start() the muxer here,
        // because our MediaFormat doesn't have the Magic Goodies.  These can only be
        // obtained from the encoder after it has started processing data.

		mMuxer = new MediaMuxer(mPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);

	    if (mMuxer == null) {
	        throw new RuntimeException("MediaMuxer create failed!");
	    }
		Log.d(LOG_TAG, "[init] Leave initWriter(): path = " + mPath);
		
	}
	

	public void setVideoTrack(MediaFormat videoFormat) {
		
		Log.d(LOG_TAG, "[setVideoTrack] Enter setVideoTrack()");
		
		if (mMuxer == null) {
			throw new IllegalArgumentException("[setVideoTrack] MediaMuxer hasn't been created");
		}
		
        if (videoFormat == null) {
            throw new IllegalArgumentException("[setVideoTrack] videoFormat must not be null");
        }
		
		// add a video track.
        mVideoTrackIndex = mMuxer.addTrack(videoFormat);
	}

	public void setAudioTrack(MediaFormat audioFormat) {
		
		Log.d(LOG_TAG, "[setVideoTrack] Enter setAudioTrack()");
		
		if (mMuxer == null) {
			throw new IllegalArgumentException("[setAudioTrack] MediaMuxer hasn't been created");
		}
		
		if (audioFormat == null) {
			throw new IllegalArgumentException("[setAudioTrack] audioFormat must not be null");
		}
		
		// add an audio track.
		mAudioTrackIndex = mMuxer.addTrack(audioFormat);
	}

	public boolean start() {
	    
	    Log.d(LOG_TAG, "[start] Enter start()");
	    
	    if (mStarted == false) {
	        mMuxer.start();
	        mStarted = true;
	        return true;
	    }
	    return false;
	}

	public void writeData(int track, ByteBuffer byteBuf, BufferInfo bufferInfo) {
		
	    String strTrack;
	    if (track == TRACK_AUDIO)
	        strTrack = "audio";
	    else
	        strTrack = "video";
	    
		Log.d(LOG_TAG, "[writeData] Enter writeData(): track = " + strTrack + 
		        " length = " + bufferInfo.size + " pts = " + bufferInfo.presentationTimeUs);
		
		if (mMuxer == null) {
			throw new IllegalArgumentException("[writeData] MediaMuxer hasn't been created");
		}
		
		if (mStarted == false) {
		    throw new IllegalStateException("[writeData] MediaMuxer hasn't been started!!");
		}
		
		if (track == TRACK_VIDEO) {
		    if (bufferInfo.flags != 0) {
		        Log.i(LOG_TAG, "[writeData] write a sync frame: track = " + strTrack);
		    }
			mMuxer.writeSampleData(mVideoTrackIndex, byteBuf, bufferInfo);
		} else {
			mMuxer.writeSampleData(mAudioTrackIndex, byteBuf, bufferInfo);
		}
		Log.d(LOG_TAG, "[writeData] Leave writeData(): track = " + strTrack);
	}
	

	public boolean stop() {
	    
	    Log.d(LOG_TAG, "[stop] Enter stop()");
	    
	    if (mStarted == true) {
	        Log.d(LOG_TAG, "[stop] stop muxer now");
	        mMuxer.stop();
	        mStarted = false;
	        
	        return true;
	    }
	    
	    return false;
	}


	public void uninit() {
		
		Log.d(LOG_TAG, "[uninitWriter] Enter uninitWriter()");
		
		if (mMuxer != null) {

			mMuxer.release();
			mMuxer = null;
		}
	}
}
